"""
项目：Ai-Engineer自适应工程设计自动化插件
作者：刘昊轩、李佳音、刘伟伟、杨明久（排序不分先后）
©版权保护
"""
# deepseek_diagnose.py
import os
import logging
from typing import Dict, Any
from dotenv import load_dotenv
from ai_integration.deepseek_client import OfficialDeepSeekClient, RateLimitError
from ai_integration.deepseek_client_ollama import OllamaClient, ModelNotFoundError
from ai_integration.fallback_client import FallbackAIClient

# 初始化日志
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger("Diagnose")

class DeepSeekDiagnoser:
    def __init__(self):
        self.tests = [
            self._test_env_vars,
            self._test_api_connection,
            self._test_model_listing,
            self._test_usage_check,
            self._test_full_cycle
        ]

    def run_diagnose(self):
        """运行全套诊断测试"""
        logger.info("🛠️ 开始DeepSeek官方API诊断...\n")
        results = {}

        for test in self.tests:
            test_name = test.__name__[6:]  # 去掉_test前缀
            try:
                result = test()
                status = "✅ 通过" if result else "❌ 失败"
                results[test_name] = {"status": status, "details": result}
                logger.info(f"{status} | {test_name}")
            except Exception as e:
                results[test_name] = {"status": "❌ 异常", "details": str(e)}
                logger.error(f"❌ 异常 | {test_name} - {str(e)}")

        logger.info("\n诊断完成")
        return results

    def _test_env_vars(self) -> Dict[str, bool]:
        """测试环境变量配置"""
        required_vars = {
            "DEEPSEEK_OFFICIAL_KEY": bool(os.getenv("DEEPSEEK_OFFICIAL_KEY"))
        }

        if not all(required_vars.values()):
            missing = [k for k, v in required_vars.items() if not v]
            raise ValueError(f"缺失环境变量: {missing}")

        return required_vars

    def _test_api_connection(self) -> Dict[str, Any]:
        """测试API连接"""
        client = OfficialDeepSeekClient()
        try:
            response = client.generate("测试连接")
            return {
                "status": "success",
                "response": response[:100] + "..." if len(response) > 100 else response
            }
        except RateLimitError:
            return {"status": "rate_limited"}
        except Exception as e:
            raise RuntimeError(f"API连接失败: {str(e)}")

    def _test_model_listing(self) -> Dict[str, Any]:
        """测试模型列表获取"""
        client = OfficialDeepSeekClient()
        try:
            models = client.list_models()
            return {
                "status": "success",
                "models": [m["id"] for m in models.get("data", [])]
            }
        except Exception as e:
            raise RuntimeError(f"获取模型列表失败: {str(e)}")

    def _test_usage_check(self) -> Dict[str, Any]:
        """测试使用情况查询"""
        client = OfficialDeepSeekClient()
        try:
            usage = client.get_usage()
            return {
                "status": "success",
                "usage": usage
            }
        except Exception as e:
            raise RuntimeError(f"获取使用情况失败: {str(e)}")

    def _test_full_cycle(self) -> Dict[str, Any]:
        """测试完整请求周期"""
        client = OfficialDeepSeekClient()
        prompt = "将以下指令转换为CAD命令：绘制直径50mm的圆，中心在原点"

        try:
            response = client.generate(prompt)
            return {
                "status": "success",
                "prompt": prompt,
                "response": response
            }
        except Exception as e:
            raise RuntimeError(f"完整请求失败: {str(e)}")

    def _test_ollama_connection(self) -> Dict[str, Any]:
        """测试Ollama服务连接"""
        try:
            client = OllamaClient()
            response = client.generate("测试连接")
            return {
                "status": "success",
                "response": response[:100] + "..." if len(response) > 100 else response
            }
        except ModelNotFoundError as e:
            return {"status": "model_not_found", "detail": str(e)}
        except Exception as e:
            raise RuntimeError(f"主服务连接失败: {str(e)}")

    def _test_model_loaded(self) -> Dict[str, Any]:
        """测试模型加载状态"""
        try:
            client = OllamaClient()
            return {"status": "loaded", "model": client.model_name}
        except Exception as e:
            raise RuntimeError(f"模型验证失败: {str(e)}")

    def _test_fallback_mechanism(self) -> Dict[str, Any]:
        """测试故障转移机制"""
        client = FallbackAIClient()

        # 模拟主服务失败
        original_endpoint = client.primary_client.base_url
        client.primary_client.base_url = "http://invalid.endpoint"

        try:
            response = client.generate("测试故障转移")
            return {
                "status": "fallback_activated",
                "response": response[:100] + "..." if len(response) > 100 else response
            }
        finally:
            client.primary_client.base_url = original_endpoint

def print_summary(results: Dict[str, Dict[str, Any]]):
    """打印诊断结果摘要"""
    print("\n诊断结果摘要：")
    for test_name, result in results.items():
        print(f"{result['status']} | {test_name}")
        if "details" in result:
            print(f"   详情: {result['details']}")

if __name__ == "__main__":
    load_dotenv()
    diagnoser = DeepSeekDiagnoser()
    results = diagnoser.run_diagnose()
    print_summary(results)